from dotenv import load_dotenv
load_dotenv()  # 加载.env文件
import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

import requests
import json
import logging
import re
import time
import hashlib
from datetime import datetime
from flask import Flask, request, jsonify
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
import torch
import threading

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('JavaAuditModel')

app = Flask(__name__)

# 缓存类定义
class AnalysisCache:
    def __init__(self):
        self.cache = {}
        self.CACHE_TIMEOUT = 3600  # 1小时缓存
    
    def get_cache_key(self, java_code):
        """生成缓存键"""
        return hashlib.md5(java_code.encode('utf-8')).hexdigest()
    
    def get(self, java_code):
        """获取缓存结果"""
        key = self.get_cache_key(java_code)
        if key in self.cache:
            cache_data = self.cache[key]
            # 检查是否过期
            if time.time() - cache_data['timestamp'] < self.CACHE_TIMEOUT:
                logger.info(f"缓存命中: {key[:8]}...")
                return cache_data['result']
        return None
    
    def set(self, java_code, result):
        """设置缓存"""
        key = self.get_cache_key(java_code)
        self.cache[key] = {
            'result': result,
            'timestamp': time.time()
        }
        logger.info(f"缓存设置: {key[:8]}...")

# 初始化缓存
analysis_cache = AnalysisCache()

# DeepSeek分析器类
class DeepSeekAnalyzer:
    def __init__(self, api_key):
        self.api_key = api_key
        self.base_url = "https://api.deepseek.com/v1/chat/completions"
        self.enabled = bool(api_key)
        logger.info(f"DeepSeek分析器初始化: {'已启用' if self.enabled else '未启用'}")
    
    def analyze(self, java_code):
        """使用DeepSeek API分析Java代码（带缓存）"""
        if not self.enabled:
            return None
        
        # 先检查缓存
        cached_result = analysis_cache.get(java_code)
        if cached_result:
            return cached_result
                
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }
        
        # 优化的提示词 - 更简洁，减少token使用
        system_prompt = """你是一个Java安全专家。分析代码漏洞，只输出JSON：
{
    "vulnerabilities": [{"type":"漏洞类型","location":"位置","description":"描述","suggestion":"建议"}],
    "secure_practices": ["安全实践"],
    "risk_assessment": "风险等级",
    "confidence": 0.9
}
只输出JSON，不要其他内容。"""
        
        # 限制代码长度，减少token消耗
        truncated_code = java_code[:2000]  # 限制2000字符
        
        data = {
            "model": "deepseek-chat",
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"分析Java代码安全：\n```java\n{truncated_code}\n```"}
            ],
            "temperature": 0.1,
            "max_tokens": 1000  # 减少输出长度
        }
        
        try:
            logger.info("调用DeepSeek API...")
            start_time = time.time()
            response = requests.post(self.base_url, headers=headers, json=data, timeout=15)  # 缩短超时
            response.raise_for_status()
            result = response.json()
            
            api_time = time.time() - start_time
            logger.info(f"DeepSeek API响应时间: {api_time:.2f}s")
            
            ai_response = result['choices'][0]['message']['content'].strip()
            
            # 清理响应
            cleaned_response = ai_response.replace('```json', '').replace('```', '').strip()
            parsed_result = json.loads(cleaned_response)
            
            # 缓存结果
            analysis_cache.set(java_code, parsed_result)
            
            return parsed_result
            
        except Exception as e:
            logger.error(f"DeepSeek API调用失败: {e}")
            return None

# 初始化DeepSeek分析器（API Key从环境变量获取）
deepseek_analyzer = DeepSeekAnalyzer(api_key=os.environ.get('DEEPSEEK_API_KEY'))

# 调试信息
print(f"调试: DEEPSEEK_API_KEY = {os.environ.get('DEEPSEEK_API_KEY')}")
print(f"调试: deepseek_analyzer.enabled = {deepseek_analyzer.enabled}")

# 服务状态统计
service_stats = {
    'start_time': datetime.now(),
    'total_requests': 0,
    'vulnerabilities_detected': 0,
    'avg_response_time': 0
}

# 漏洞检测规则库
VULNERABILITY_PATTERNS = {
    'SQL注入': [
        r"Statement\.execute\s*\(",
        r"executeQuery\s*\(",
        r"executeUpdate\s*\(",
        r"\.createStatement\s*\(",
        r"\"SELECT.*\"\s*\+",
        r"\"INSERT.*\"\s*\+",
        r"\"UPDATE.*\"\s*\+",
        r"\"DELETE.*\"\s*\+"
    ],
    '命令执行': [
        r"Runtime\.getRuntime\s*\(\s*\)\.exec\s*\(",
        r"ProcessBuilder\s*\(",
        r"\.exec\s*\("
    ],
    '路径遍历': [
        r"File\s*\(\s*[^)]*\s*\+\s*[^)]*\)",
        r"new FileInputStream\s*\(\s*[^)]*\s*\+\s*[^)]*\)",
        r"\.getAbsolutePath\s*\(\s*\)"
    ],
    'XSS跨站脚本': [
        r"response\.getWriter\s*\(\s*\)\.print\s*\(",
        r"response\.getWriter\s*\(\s*\)\.write\s*\(",
        r"PrintWriter.*\.print\s*\("
    ],
    '不安全的反序列化': [
        r"readObject\s*\(",
        r"ObjectInputStream\s*\(",
        r"\.readUnshared\s*\("
    ]
}

# 安全代码模式
SAFE_PATTERNS = [
    r"PreparedStatement",
    r"usePreparedStatements",
    r"参数化查询",
    r"inputValidation",
    r"encodeForHTML",
    r"StringUtils\.escape"
]

# AI模型初始化
def init_ai_model():
    """初始化AI模型"""
    try:
        model = pipeline(
            "text-classification",
            model="microsoft/codebert-base",
            tokenizer="microsoft/codebert-base",
            framework="pt"
        )
        logger.info("✅ AI模型加载成功")
        return model
    except Exception as e:
        logger.warning(f"❌ AI模型加载失败: {e}")
        return None

# 全局AI模型实例
ai_model = init_ai_model()
AI_ENABLED = ai_model is not None

def rule_based_detection(java_code):
    """基于规则的漏洞检测"""
    # 先检查安全模式（优先级更高）
    safe_indicators = [
        r"PreparedStatement",
        r"stmt\.setString\s*\(",
        r"stmt\.setInt\s*\(",
        r"stmt\.setBoolean\s*\(",
        r"stmt\.setDate\s*\(",
        r"参数化查询",
        r"StringUtils\.escape",
        r"encodeForHTML"
    ]
    
    safe_count = 0
    for pattern in safe_indicators:
        if re.search(pattern, java_code, re.IGNORECASE):
            safe_count += 1
    
    # 如果有明确的安全模式，直接返回安全
    if safe_count >= 2:
        return {
            "has_vulnerability": False,
            "vulnerability_type": "安全代码（参数化查询）",
            "confidence": min(0.1 + (safe_count * 0.2), 0.9),
            "location": "无",
            "risk_level": "安全",
            "detection_method": "规则引擎",
            "safe_patterns_found": safe_count
        }
    
    # 原有的漏洞检测逻辑
    vulnerabilities = []
    
    # 检查每种漏洞类型
    for vuln_type, patterns in VULNERABILITY_PATTERNS.items():
        for pattern in patterns:
            if re.search(pattern, java_code, re.IGNORECASE):
                vulnerabilities.append(vuln_type)
                break
    
    # 检查安全代码模式
    safe_count = sum(1 for pattern in SAFE_PATTERNS if re.search(pattern, java_code, re.IGNORECASE))
    
    # 构建结果
    if vulnerabilities:
        base_confidence = 0.7 + (len(vulnerabilities) * 0.05)
        confidence = min(base_confidence, 0.95)
        
        return {
            "has_vulnerability": True,
            "vulnerability_type": "、".join(vulnerabilities),
            "confidence": round(confidence, 2),
            "location": "需AST解析精确定位",
            "risk_level": "高危" if len(vulnerabilities) > 1 else "中危",
            "detection_method": "规则引擎",
            "safe_patterns_found": safe_count
        }
    else:
        safe_confidence = 0.1 + (safe_count * 0.1)
        return {
            "has_vulnerability": False,
            "vulnerability_type": "未发现漏洞",
            "confidence": round(min(safe_confidence, 0.5), 2),
            "location": "无",
            "risk_level": "安全",
            "detection_method": "规则引擎",
            "safe_patterns_found": safe_count
        }

def ai_enhanced_detection(java_code, rule_result):
    """AI增强的漏洞检测 - 使用DeepSeek API"""
    if not deepseek_analyzer.enabled:
        logger.warning("DeepSeek API未启用，使用规则引擎结果")
        return rule_result
    
    try:
        # 调用DeepSeek API
        ai_result = deepseek_analyzer.analyze(java_code)
        
        if not ai_result:
            logger.warning("DeepSeek API返回空结果")
            return rule_result
        
        logger.info(f"DeepSeek分析成功: {ai_result.get('risk_assessment', '未知')}")
        
        # 解析DeepSeek结果并与规则结果融合
        return merge_ai_rule_results(ai_result, rule_result, java_code)
        
    except Exception as e:
        logger.error(f"DeepSeek分析失败: {e}")
        return rule_result

def merge_ai_rule_results(ai_result, rule_result, java_code):
    """融合AI结果和规则引擎结果"""
    
    # 如果AI检测到漏洞
    ai_vulnerabilities = ai_result.get('vulnerabilities', [])
    ai_secure_practices = ai_result.get('secure_practices', [])
    
    # 情况1: 规则引擎说安全，但AI发现漏洞
    if not rule_result["has_vulnerability"] and ai_vulnerabilities:
        vuln_types = list(set([vuln.get('type', '潜在风险') for vuln in ai_vulnerabilities]))
        
        return {
            "has_vulnerability": True,
            "vulnerability_type": "、".join(vuln_types),
            "confidence": min(ai_result.get('confidence', 0.7) + 0.1, 0.95),
            "location": "；".join([vuln.get('location', '未知位置') for vuln in ai_vulnerabilities[:2]]),
            "risk_level": "高危" if len(vuln_types) > 1 else "中危",
            "detection_method": "DeepSeek AI",
            "ai_confidence": ai_result.get('confidence', 0.7),
            "ai_details": ai_vulnerabilities,
            "secure_practices": ai_secure_practices
        }
    
    # 情况2: 规则引擎和AI都发现漏洞
    elif rule_result["has_vulnerability"] and ai_vulnerabilities:
        # 提升置信度，合并漏洞类型
        rule_vulns = set(rule_result["vulnerability_type"].split("、"))
        ai_vulns = set([vuln.get('type', '') for vuln in ai_vulnerabilities])
        all_vulns = list(rule_vulns.union(ai_vulns))
        all_vulns = [v for v in all_vulns if v]
        
        rule_result.update({
            "vulnerability_type": "、".join(all_vulns),
            "confidence": min(rule_result["confidence"] + 0.15, 0.98),
            "detection_method": "规则+DeepSeek双确认",
            "ai_confidence": ai_result.get('confidence', 0.7),
            "ai_details": ai_vulnerabilities,
            "secure_practices": ai_secure_practices
        })
        return rule_result
    
    # 情况3: 两者都认为安全，但AI识别出安全实践
    elif not rule_result["has_vulnerability"] and ai_secure_practices:
        rule_result.update({
            "secure_practices": ai_secure_practices,
            "detection_method": "规则+DeepSeek双重确认",
            "confidence": min(rule_result["confidence"] + 0.1, 0.8)
        })
        return rule_result
    
    # 其他情况返回规则结果
    return rule_result

def update_stats(response_time, has_vulnerability):
    """更新服务统计信息"""
    service_stats['total_requests'] += 1
    if has_vulnerability:
        service_stats['vulnerabilities_detected'] += 1
    
    # 更新平均响应时间（移动平均）
    old_avg = service_stats['avg_response_time']
    total_reqs = service_stats['total_requests']
    service_stats['avg_response_time'] = (old_avg * (total_reqs - 1) + response_time) / total_reqs

# ==================== API路由 ====================

@app.route('/health', methods=['GET'])
def health_check():
    """健康检查接口"""
    return jsonify({
        "status": "running",
        "service": "Java代码审计模型",
        "version": "3.0",
        "ai_engine": "DeepSeek API" if deepseek_analyzer.enabled else "未启用",
        "rule_engine": "已启用",
        "deepseek_enabled": deepseek_analyzer.enabled,
        "uptime": str(datetime.now() - service_stats['start_time']).split('.')[0]
    })

@app.route('/cache/stats', methods=['GET'])
def cache_stats():
    """缓存统计信息"""
    return jsonify({
        "cache_size": len(analysis_cache.cache),
        "cache_timeout": analysis_cache.CACHE_TIMEOUT,
        "cache_keys": list(analysis_cache.cache.keys())[:5]  # 显示前5个键
    })

@app.route('/analyze', methods=['POST'])
def analyze_code():
    """代码分析接口"""
    start_time = time.time()
    
    try:
        data = request.get_json()
        if not data:
            return jsonify({"error": "请求体必须为JSON格式"}), 400
        
        java_code = data.get('code', '').strip()
        if not java_code:
            return jsonify({"error": "代码内容不能为空"}), 400
        
        if len(java_code) > 10000:
            return jsonify({"error": "代码长度超过限制(10000字符)"}), 400
        
        logger.info(f"分析代码请求, 长度: {len(java_code)}")
        
        # 1. 规则引擎检测
        rule_result = rule_based_detection(java_code)
        
        # 2. AI增强分析
        final_result = ai_enhanced_detection(java_code, rule_result)
        
        # 3. 更新统计
        response_time = time.time() - start_time
        update_stats(response_time, final_result["has_vulnerability"])
        
        # 4. 添加元数据
        final_result.update({
            "response_time": round(response_time, 3),
            "code_length": len(java_code),
            "timestamp": datetime.now().isoformat()
        })
        
        logger.info(f"分析完成: 漏洞={final_result['has_vulnerability']}, 耗时={response_time:.3f}s")
        
        return jsonify(final_result)
        
    except Exception as e:
        logger.error(f"分析过程出错: {e}")
        return jsonify({"error": "服务器内部错误", "details": str(e)}), 500

@app.route('/stats', methods=['GET'])
def get_stats():
    """获取服务统计信息"""
    stats = service_stats.copy()
    stats['current_time'] = datetime.now().isoformat()
    stats['uptime'] = str(datetime.now() - stats['start_time']).split('.')[0]
    return jsonify(stats)

@app.route('/patterns', methods=['GET'])
def get_patterns():
    """获取支持的检测模式"""
    return jsonify({
        "vulnerability_patterns": list(VULNERABILITY_PATTERNS.keys()),
        "safe_patterns": SAFE_PATTERNS,
        "total_rules": sum(len(patterns) for patterns in VULNERABILITY_PATTERNS.values()) + len(SAFE_PATTERNS)
    })

@app.route('/')
def home():
    """根路径"""
    return jsonify({
        "message": "Transformer-based Java代码审计模型服务",
        "version": "3.0",
        "endpoints": {
            "健康检查": "GET /health",
            "代码分析": "POST /analyze",
            "服务统计": "GET /stats",
            "检测模式": "GET /patterns",
            "缓存统计": "GET /cache/stats"
        },
        "ai_engines": {
            "deepseek": deepseek_analyzer.enabled,
            "local_model": AI_ENABLED
        }
    })

# 请求钩子
@app.before_request
def before_request():
    """请求前处理"""
    request.start_time = time.time()

@app.after_request
def after_request(response):
    """请求后处理"""
    if hasattr(request, 'start_time'):
        response_time = time.time() - request.start_time
        logger.info(f"请求 {request.path} - 状态 {response.status_code} - 耗时 {response_time:.3f}s")
    
    # 添加CORS头
    response.headers.add('Access-Control-Allow-Origin', '*')
    response.headers.add('Access-Control-Allow-Headers', 'Content-Type')
    response.headers.add('Access-Control-Allow-Methods', 'GET, POST, OPTIONS')
    
    return response

@app.errorhandler(404)
def not_found(error):
    return jsonify({"error": "接口不存在"}), 404

@app.errorhandler(500)
def internal_error(error):
    return jsonify({"error": "服务器内部错误"}), 500

if __name__ == '__main__':
    logger.info("🚀 启动Java代码审计模型服务...")
    logger.info(f"📡 服务地址: http://127.0.0.1:5000")
    logger.info(f"🤖 DeepSeek API状态: {'已启用' if deepseek_analyzer.enabled else '未启用'}")
    logger.info(f"🤖 本地AI模型状态: {'已启用' if AI_ENABLED else '未启用'}")
    logger.info(f"💾 缓存系统: 已启用")
    
    app.run(host='0.0.0.0', port=5000, debug=False)